/*  NetBare - An android network capture and injection library.
 *  Copyright (C) 2018-2019 Megatron King
 *  Copyright (C) 2018-2019 GuoShi
 *
 *  NetBare is free software: you can redistribute it and/or modify it under the terms
 *  of the GNU General Public License as published by the Free Software Found-
 *  ation, either version 3 of the License, or (at your option) any later version.
 *
 *  NetBare is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 *  without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
 *  PURPOSE. See the GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along with NetBare.
 *  If not, see <http://www.gnu.org/licenses/>.
 */
package com.github.megatronking.netbare.ssl;

import android.os.Build;

import com.github.megatronking.netbare.NetBareUtils;

import org.bouncycastle.asn1.ASN1EncodableVector;
import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERSequence;
import org.bouncycastle.asn1.x500.X500Name;
import org.bouncycastle.asn1.x500.X500NameBuilder;
import org.bouncycastle.asn1.x500.style.BCStyle;
import org.bouncycastle.asn1.x509.BasicConstraints;
import org.bouncycastle.asn1.x509.Extension;
import org.bouncycastle.asn1.x509.GeneralName;
import org.bouncycastle.asn1.x509.KeyPurposeId;
import org.bouncycastle.asn1.x509.KeyUsage;
import org.bouncycastle.asn1.x509.SubjectKeyIdentifier;
import org.bouncycastle.asn1.x509.SubjectPublicKeyInfo;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.bc.BcX509ExtensionUtils;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.OperatorCreationException;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.NoSuchProviderException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.SignatureException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.util.Date;
import java.util.Random;

/**
 * Generates self-signed certificate by {@link JKS}.
 *
 * @author Megatron King
 * @since 2018-11-08 02:23
 */
public final class CertificateGenerator {

    private static final String KEY_STORE_TYPE = "PKCS12";
    private static final String KEYGEN_ALGORITHM = "RSA";
    private static final String SECURE_RANDOM_ALGORITHM = "SHA1PRNG";

    private static final String PROVIDER_NAME = BouncyCastleProvider.PROVIDER_NAME;

    private static final int ROOT_KEY_SIZE = 2048;
    private static final int SERVER_KEY_SIZE = 1024;

    /**
     * The signature algorithm starting with the message digest to use when signing certificates.
     * On 64-bit systems this should be set to SHA512, on 32-bit systems this is SHA256. On 64-bit
     * systems, SHA512 generally performs better than SHA256; see this question for details:
     * http://crypto.stackexchange.com/questions/26336/sha512-faster-than-sha256
     */
    private static final String SIGNATURE_ALGORITHM = (is32BitJvm() ? "SHA256" : "SHA512") +
            "WithRSAEncryption";

    /**
     * The milliseconds of 1 day.
     */
    private static final long ONE_DAY = 86400000L;

    /**
     * Current time minus 1 year, just in case software clock goes back due to time synchronization.
     */
    private static final Date NOT_BEFORE = new Date(System.currentTimeMillis() - ONE_DAY * 365);

    /**
     * The maximum possible value in X.509 specification: 9999-12-31 23:59:59,
     * new Date(253402300799000L), but Apple iOS 8 fails with a certificate
     * expiration date grater than Mon, 24 Jan 6084 02:07:59 GMT (issue #6).
     *
     * Hundred years in the future from starting the proxy should be enough.
     */
    private static final Date NOT_AFTER = new Date(System.currentTimeMillis() + ONE_DAY * 365 * 10);

    /**
     * Generate a root keystore by a given {@link JKS}.
     *
     * @param jks A java keystore object.
     * @return A root {@link KeyStore}.
     */
    public KeyStore generateRoot(JKS jks)
            throws KeyStoreException, CertificateException, NoSuchAlgorithmException,
            IOException, OperatorCreationException {
        KeyPair keyPair = generateKeyPair(ROOT_KEY_SIZE);

        X500NameBuilder nameBuilder = new X500NameBuilder(BCStyle.INSTANCE);
        nameBuilder.addRDN(BCStyle.CN, jks.commonName());
        nameBuilder.addRDN(BCStyle.O, jks.organization());
        nameBuilder.addRDN(BCStyle.OU, jks.organizationalUnitName());
        X500Name issuer = nameBuilder.build();

        PublicKey pubKey = keyPair.getPublic();

        X509v3CertificateBuilder generator = new JcaX509v3CertificateBuilder(
                issuer, BigInteger.valueOf(randomSerial()), NOT_BEFORE, NOT_AFTER, issuer, pubKey);
        generator.addExtension(Extension.subjectKeyIdentifier, false,
                createSubjectKeyIdentifier(pubKey));
        generator.addExtension(Extension.basicConstraints, true,
                new BasicConstraints(true));

        KeyUsage usage = new KeyUsage(KeyUsage.keyCertSign | KeyUsage.digitalSignature |
                KeyUsage.keyEncipherment | KeyUsage.dataEncipherment | KeyUsage.cRLSign);
        generator.addExtension(Extension.keyUsage, false, usage);

        ASN1EncodableVector purposes = new ASN1EncodableVector();
        purposes.add(KeyPurposeId.id_kp_serverAuth);
        purposes.add(KeyPurposeId.id_kp_clientAuth);
        purposes.add(KeyPurposeId.anyExtendedKeyUsage);
        generator.addExtension(Extension.extendedKeyUsage, false,
                new DERSequence(purposes));

        X509Certificate cert = signCertificate(generator, keyPair.getPrivate());

        KeyStore result = KeyStore.getInstance(KEY_STORE_TYPE);
        result.load(null, null);
        result.setKeyEntry(jks.alias(), keyPair.getPrivate(), jks.password(),
                new Certificate[] { cert });
        return result;
    }

    public KeyStore generateServer(String commonName, JKS jks,
                                          Certificate caCert, PrivateKey caPrivKey)
            throws NoSuchAlgorithmException, NoSuchProviderException,
            IOException, OperatorCreationException, CertificateException,
            InvalidKeyException, SignatureException, KeyStoreException {

        KeyPair keyPair = generateKeyPair(SERVER_KEY_SIZE);

        X500Name issuer = new X509CertificateHolder(caCert.getEncoded()).getSubject();
        BigInteger serial = BigInteger.valueOf(randomSerial());
        X500NameBuilder name = new X500NameBuilder(BCStyle.INSTANCE);
        name.addRDN(BCStyle.CN, commonName);
        name.addRDN(BCStyle.O, jks.certOrganisation());
        name.addRDN(BCStyle.OU, jks.certOrganizationalUnitName());
        X500Name subject = name.build();

        X509v3CertificateBuilder builder = new JcaX509v3CertificateBuilder(issuer, serial, NOT_BEFORE,
                new Date(System.currentTimeMillis() + ONE_DAY), subject, keyPair.getPublic());
        builder.addExtension(Extension.subjectKeyIdentifier, false,
                createSubjectKeyIdentifier(keyPair.getPublic()));
        builder.addExtension(Extension.basicConstraints, false,
                new BasicConstraints(false));
        builder.addExtension(Extension.subjectAlternativeName, false,
                new DERSequence(new GeneralName(GeneralName.dNSName, commonName)));

        X509Certificate cert = signCertificate(builder, caPrivKey);

        cert.checkValidity(new Date());
        cert.verify(caCert.getPublicKey());

        KeyStore result = KeyStore.getInstance(KeyStore.getDefaultType());
        result.load(null, null);
        Certificate[] chain = { cert, caCert };
        result.setKeyEntry(jks.alias(), keyPair.getPrivate(), jks.password(), chain);
        return result;
    }

    public String keyStoreType() {
        return KEY_STORE_TYPE;
    }

    private KeyPair generateKeyPair(int keySize) throws NoSuchAlgorithmException {
        KeyPairGenerator generator = KeyPairGenerator.getInstance(KEYGEN_ALGORITHM);
        SecureRandom secureRandom = SecureRandom.getInstance(SECURE_RANDOM_ALGORITHM);
        generator.initialize(keySize, secureRandom);
        return generator.generateKeyPair();
    }

    private long randomSerial() {
        final Random rnd = new Random();
        rnd.setSeed(System.currentTimeMillis());
        // prevent browser certificate caches, cause of doubled serial numbers
        // using 48bit random number
        long sl = ((long) rnd.nextInt()) << 32 | (rnd.nextInt() & 0xFFFFFFFFL);
        // let reserve of 16 bit for increasing, serials have to be positive
        sl = sl & 0x0000FFFFFFFFFFFFL;
        return sl;
    }

    private static SubjectKeyIdentifier createSubjectKeyIdentifier(Key key) throws IOException {
        ByteArrayInputStream bIn = new ByteArrayInputStream(key.getEncoded());
        ASN1InputStream is = null;
        try {
            is = new ASN1InputStream(bIn);
            ASN1Sequence seq = (ASN1Sequence) is.readObject();
            SubjectPublicKeyInfo info = SubjectPublicKeyInfo.getInstance(seq);
            return new BcX509ExtensionUtils().createSubjectKeyIdentifier(info);
        } finally {
            NetBareUtils.closeQuietly(is);
        }
    }

    private static X509Certificate signCertificate(X509v3CertificateBuilder certificateBuilder,
            PrivateKey signedWithPrivateKey) throws OperatorCreationException,
            CertificateException {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.P) {
            ContentSigner signer = new JcaContentSignerBuilder(SIGNATURE_ALGORITHM)
                    .build(signedWithPrivateKey);
            return new JcaX509CertificateConverter()
                    .getCertificate(certificateBuilder.build(signer));
        } else {
            ContentSigner signer = new JcaContentSignerBuilder(SIGNATURE_ALGORITHM)
                    .setProvider(PROVIDER_NAME)
                    .build(signedWithPrivateKey);
            return new JcaX509CertificateConverter()
                    .setProvider(PROVIDER_NAME)
                    .getCertificate(certificateBuilder.build(signer));
        }
    }

    /**
     * Uses the non-portable system property sun.arch.data.model to help
     * determine if we are running on a 32-bit JVM. Since the majority of modern
     * systems are 64 bits, this method "assumes" 64 bits and only returns true
     * if sun.arch.data.model explicitly indicates a 32-bit JVM.
     *
     * @return true if we can determine definitively that this is a 32-bit JVM,
     *         otherwise false
     */
    private static boolean is32BitJvm() {
        Integer bits = Integer.getInteger("sun.arch.data.model");
        return bits != null && bits == 32;
    }

}
